# 导入必要的库
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import jieba
import joblib


# ===========================
# 1️⃣ 数据加载
# ===========================
# 读取训练集和测试集
train_data = pd.read_csv(r'C:\Users\Xander.Zheng\Desktop\train_data.csv')
test_data = pd.read_csv(r'C:\Users\Xander.Zheng\Desktop\test_data.csv')

# 查看数据格式
print("训练集数据：")
print(train_data.head())
print("\n测试集数据：")
print(test_data.head())

# ===========================
# 2️⃣ 数据预处理
# ===========================
# 合并标题和正文，形成文本特征
train_data['text'] = train_data['标题'].fillna('') + ' ' + train_data['正文'].fillna('')
test_data['text'] = test_data['标题'].fillna('') + ' ' + test_data['正文'].fillna('')

# 提取文本和标签
X_train_text = train_data['text']
y_train = train_data['正负面']
X_test_text = test_data['text']
y_test = test_data['正负面']

# ===========================
# 3️⃣ 文本分词
# ===========================
def chinese_tokenizer(text):
    return jieba.lcut(text)

# ===========================
# 4️⃣ 特征提取 - TF-IDF
# ===========================
tfidf_vectorizer = TfidfVectorizer(tokenizer=chinese_tokenizer, max_features=10000)

# 拟合训练集并转换成 TF-IDF 特征向量
X_train_tfidf = tfidf_vectorizer.fit_transform(X_train_text)
X_test_tfidf = tfidf_vectorizer.transform(X_test_text)

print(f"训练集特征维度: {X_train_tfidf.shape}")
print(f"测试集特征维度: {X_test_tfidf.shape}")

# ===========================
# 5️⃣ 模型训练 - SVM
# ===========================
# 初始化 SVM 分类器
svm_model = SVC(kernel='linear', C=1.0, probability=True, random_state=42)

# 训练模型
svm_model.fit(X_train_tfidf, y_train)

# ===========================
# 6️⃣ 模型评估
# ===========================
# 预测测试集情感
y_pred = svm_model.predict(X_test_tfidf)

# 输出模型评估指标
print("\n🎯 模型评估结果：")
print(f"准确率: {accuracy_score(y_test, y_pred):.4f}")
print("\n分类报告:")
print(classification_report(y_test, y_pred))
print("\n混淆矩阵:")
print(confusion_matrix(y_test, y_pred))

# ===========================
# 7️⃣ 单条新闻情感预测
# ===========================
def predict_sentiment(text):
    text_tfidf = tfidf_vectorizer.transform([text])
    sentiment = svm_model.predict(text_tfidf)[0]
    return sentiment

# 测试单条新闻
sample_text = ("哈哈哈")
print(f"\n新闻: {sample_text}")
print(f"情感预测: {predict_sentiment(sample_text) }")
print(f"情感预测: {'正面' if predict_sentiment(sample_text) == 1 else '负面'}")
# 保存模型
joblib.dump(svm_model, 'svm_model.joblib')
# 保存 TfidfVectorizer 对象
joblib.dump(tfidf_vectorizer, 'tfidf_vectorizer.joblib')
